{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1622e1563a35aa32",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": "# Conversational RAG with Burr and Hamilton\n\nSee [GitHub example](https://github.com/apache/burr/tree/main/examples/conversational-rag/simple_example) and the accompanying video walkthrough:\n\n<a href=\"https://www.youtube.com/watch?v=t54DCiOH270\"><img src=\"https://img.youtube.com/vi/t54DCiOH270/maxresdefault.jpg\" alt=\"RAG Chatbot Video\" width=\"600\"></a>"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9069a73d207fd136",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "!pip install burr[start] sf-hamilton[visualization] openai "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-05-10T21:02:28.908380Z",
     "start_time": "2024-05-10T21:02:18.236814Z"
    }
   },
   "outputs": [],
   "source": [
    "# Importing the necessary libraries\n",
    "import pprint\n",
    "from typing import Tuple\n",
    "from hamilton import dataflows, driver\n",
    "import burr.core\n",
    "from burr.core import ApplicationBuilder, State, default, expr\n",
    "from burr.core.action import action\n",
    "from burr.tracking import LocalTrackingClient\n",
    "import uuid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f1d578469a918b1",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Load your \"chain\" or conversational RAG \"pipeline\"\n",
    "\n",
    "We use Hamilton here. But you could use LangChain, etc., or forgo them and write your own code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c6a018aff1154f0b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:47:42.057246Z",
     "start_time": "2024-03-26T20:47:40.500075Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Loads Hamilton DAG\n",
    "conversational_rag = dataflows.import_module(\"conversational_rag\")\n",
    "conversational_rag_driver = (\n",
    "    driver.Builder()\n",
    "    .with_config({})  # replace with configuration as appropriate\n",
    "    .with_modules(conversational_rag)\n",
    "    .build()\n",
    ")\n",
    "conversational_rag_driver.display_all_functions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82b3515afd2de6e4",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Create the actions that will constitute our application\n",
    "\n",
    "We will use the functional (vs class) approach to declaring actions here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6433dad5abc6eb16",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:47:42.065785Z",
     "start_time": "2024-03-26T20:47:42.059539Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "@action(\n",
    "    reads=[\"question\", \"chat_history\"],\n",
    "    writes=[\"chat_history\"],\n",
    ")\n",
    "def ai_converse(state: State, vector_store: object) -> Tuple[dict, State]:\n",
    "    \"\"\"AI conversing step. Uses Hamilton to execute the conversational pipeline.\"\"\"\n",
    "    result = conversational_rag_driver.execute(\n",
    "        [\"conversational_rag_response\"],\n",
    "        inputs={\n",
    "            \"question\": state[\"question\"],\n",
    "            \"chat_history\": state[\"chat_history\"],\n",
    "        },\n",
    "        # we use overrides here because we want to pass in the vector store\n",
    "        overrides={\n",
    "            \"vector_store\": vector_store,\n",
    "        }\n",
    "    )\n",
    "    new_history = f\"AI: {result['conversational_rag_response']}\"\n",
    "    return result, state.append(chat_history=new_history)\n",
    "\n",
    "\n",
    "@action(\n",
    "    reads=[],\n",
    "    writes=[\"question\", \"chat_history\"],\n",
    ")\n",
    "def human_converse(state: State, user_question: str) -> Tuple[dict, State]:\n",
    "    \"\"\"Human converse step -- make sure we get input, and store it as state.\"\"\"\n",
    "    state = state.update(question=user_question).append(chat_history=f\"Human: {user_question}\")\n",
    "    return {\"question\": user_question}, state"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4439e0dd39441414",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Add a hook to print the steps -- optional but shows that Burr is pluggable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e3c3af84fab14f39",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "from burr.core import Action\n",
    "from burr.lifecycle import PostRunStepHook, PreRunStepHook\n",
    "\n",
    "class PrintStepHook(PostRunStepHook, PreRunStepHook):\n",
    "    \"\"\"Custom hook to print the action/result after each step.\"\"\"\n",
    "\n",
    "    def pre_run_step(self, action: Action, **future_kwargs):\n",
    "        if action.name == \"ai_converse\":\n",
    "            print(\"🤔 AI is thinking...\")\n",
    "        if action.name == \"human_converse\":\n",
    "            print(\"⏳Processing input from user...\")\n",
    "\n",
    "    def post_run_step(self, *, state: \"State\", action: Action, result: dict, **future_kwargs):\n",
    "        if action.name == \"human_converse\":\n",
    "            print(\"🎙💬\", result[\"question\"], \"\\n\")\n",
    "        if action.name == \"ai_converse\":\n",
    "            print(\"🤖💬\", result[\"conversational_rag_response\"], \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9579b47aac2c53a0",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Create the application\n",
    "\n",
    "We now create the application, which is a collection of actions, and then set the transitions between the actions based on values in State.\n",
    "\n",
    "We also initialize initial values etc to populate the application with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "77e9f67b660a0953",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:50:25.642660Z",
     "start_time": "2024-03-26T20:50:25.616245Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# what we will do RAG over.\n",
    "initial_documents = [\n",
    "    \"harrison worked at kensho\",\n",
    "    \"stefan worked at Stitch Fix\",\n",
    "    \"stefan likes tacos\",\n",
    "    \"elijah worked at TwoSigma\",\n",
    "    \"elijah likes mango\",\n",
    "    \"stefan used to work at IBM\",\n",
    "    \"elijah likes to go biking\",\n",
    "    \"stefan likes to bake sourdough\",\n",
    "]\n",
    "# bootstrap the vector store;\n",
    "vector_store = conversational_rag_driver.execute(\n",
    "    [\"vector_store\"],\n",
    "    inputs={\"input_texts\": initial_documents})[\"vector_store\"]\n",
    "# what we will initialize the application with\n",
    "initial_state = {\n",
    "    \"question\": \"\",\n",
    "    \"chat_history\": [],\n",
    "}\n",
    "\n",
    "app_id = str(uuid.uuid4())\n",
    "app = (\n",
    "    ApplicationBuilder()\n",
    "    # add the actions\n",
    "    .with_actions(\n",
    "        # bind the vector store to the AI conversational step\n",
    "        ai_converse=ai_converse.bind(vector_store=vector_store),\n",
    "        human_converse=human_converse,\n",
    "        terminal=burr.core.Result(\"chat_history\"),\n",
    "    )\n",
    "    # set the transitions between actions\n",
    "    .with_transitions(\n",
    "        (\"ai_converse\", \"human_converse\", default),\n",
    "        (\"human_converse\", \"terminal\", expr(\"'exit' in question\")),\n",
    "        (\"human_converse\", \"ai_converse\", default),\n",
    "    )\n",
    "    # add identifiers that will help track the application\n",
    "    .with_identifiers(app_id=app_id, partition_key=\"sample_user\")\n",
    "    # initialize the state\n",
    "    .with_state(**initial_state)\n",
    "    # say what the initial action is\n",
    "    .with_entrypoint(\"human_converse\")\n",
    "    # add a hook to print the steps -- optional but shows that Burr is pluggable\n",
    "    .with_hooks(PrintStepHook())\n",
    "    # add tracking -- this will show up in the BURR UI.\n",
    "    .with_tracker(project=\"demo:conversational-rag\")\n",
    "    # build the application\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bf5d1e084a791fa5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:48:39.379712Z",
     "start_time": "2024-03-26T20:48:38.701659Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# let's visualize what we have\n",
    "app.visualize(include_conditions=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "430bab287b6ad9a",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Let's run the app. \n",
    "\n",
    "Let's run it a step at a time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8bcfe9ca48f87618",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:50:38.782857Z",
     "start_time": "2024-03-26T20:50:28.797204Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "app.reset_to_entrypoint() # reset the app to the entrypoint\n",
    "user_question = input(\"Ask something (or type exit to quit): \")\n",
    "previous_action, result, state = app.step(\n",
    "    inputs={\"user_question\": user_question},\n",
    ")\n",
    "print(f\"Ran action {previous_action.name} with result:\\n {pprint.pformat(result)} \\n and state:\\n {pprint.pformat(state.get_all())}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "81940578d58fd602",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:50:44.662755Z",
     "start_time": "2024-03-26T20:50:41.919782Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# now let's run the AI conversational step\n",
    "previous_action, result, state = app.step()\n",
    "print(f\"Ran action {previous_action.name} with result:\\n {pprint.pformat(result)} \\n and state:\\n {pprint.pformat(state.get_all())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36ec2f4908c2dde2",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Let's now run the app to completion\n",
    "\n",
    "You could do the above for each action. Or you could tell the app to run until certain\n",
    "actions/conditions are met."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "be6c573158b65cb1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:52:21.364028Z",
     "start_time": "2024-03-26T20:50:52.382808Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "print(f\"Running RAG with initial state:\\n {pprint.pformat(app.state.get_all())}\")\n",
    "while True:\n",
    "    user_question = input(\"Ask something (or type exit to quit): \")\n",
    "    previous_action, result, state = app.run(\n",
    "        halt_before=[\"human_converse\"],\n",
    "        halt_after=[\"terminal\"],\n",
    "        inputs={\"user_question\": user_question},\n",
    "    )\n",
    "    if previous_action.name == \"terminal\":\n",
    "        # reached the end\n",
    "        pprint.pprint(result)\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "169946a65f977df9",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "## Reloading from prior state\n",
    "\n",
    "Burr makes it easy to reload from a prior state. In this example we'll just use what is logged to the tracker to \"go back in time\" and reload the application to that state. \n",
    "\n",
    "This is useful for debugging, building the application itself, etc.\n",
    "\n",
    "There are two ways to load prior state:\n",
    "1. Load the state outside the Burr Application. i.e. pass it in as initial state.\n",
    "2. Use the ApplicationBuilder .initialize_from() method. You can optionally \"fork\" or continue an existing application this way.\n",
    "\n",
    "The difference between them, is that the first method is more flexible and gives you total control.\n",
    "The second method gives you two options: \n",
    "(a) to explicitly \"fork\" i.e. create a new application from a certain point in time from an existing application_id.\n",
    "(b) allows you to \"pick up where you left off\" with an existing application_id, e.g. in the case of a crash, or if you wanted to start from the last conversation with a user, and continue.\n",
    "\n",
    "Below we show how to do the first method. Then after that the second method, to show how to fork the conversation,\n",
    "as well as pick up the prior conversation from where it left off.  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c7f4dd64f73ed2d8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:52:37.419848Z",
     "start_time": "2024-03-26T20:52:37.393869Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# set up for rewinding to a prior state -- loading it in as initial state\n",
    "prior_app_id = app_id\n",
    "last_sequence_id = app.sequence_id\n",
    "rewind_to_sequence_id = last_sequence_id - 2\n",
    "new_app_id = str(uuid.uuid4())\n",
    "\n",
    "project_name = \"demo:conversational-rag\"\n",
    "# we use the tracking client here to get the state of the application at a prior sequence_id\n",
    "tracker = LocalTrackingClient(project=project_name)\n",
    "persisted_state = tracker.load(partition_key=\"sample_user\", \n",
    "                               app_id=prior_app_id, \n",
    "                               sequence_id=rewind_to_sequence_id)\n",
    "state_values = persisted_state['state'].get_all()\n",
    "print(f\"Loaded state from app_id:{prior_app_id}, \"\n",
    "      f\"sequence_id:{rewind_to_sequence_id}::\\n \"\n",
    "      f\"{pprint.pformat(state_values)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5ee618e3-15c0-403b-bc96-3a2faaea457e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# build our other application using the state we just loaded.\n",
    "other_app = (\n",
    "    ApplicationBuilder()\n",
    "    # add the actions\n",
    "    .with_actions(\n",
    "        # bind the vector store to the AI conversational step\n",
    "        ai_converse=ai_converse.bind(vector_store=vector_store),\n",
    "        human_converse=human_converse,\n",
    "        terminal=burr.core.Result(\"chat_history\"),\n",
    "    )\n",
    "    # set the transitions between actions\n",
    "    .with_transitions(\n",
    "        (\"ai_converse\", \"human_converse\", default),\n",
    "        (\"human_converse\", \"terminal\", expr(\"'exit' in question\")),\n",
    "        (\"human_converse\", \"ai_converse\", default),\n",
    "    )\n",
    "    # add identifiers that will help track the application\n",
    "    .with_identifiers(app_id=new_app_id, partition_key=\"sample_user\")\n",
    "    # set state to prior state\n",
    "    .with_state(**persisted_state[\"state\"].get_all())\n",
    "    # say where we want to start\n",
    "    .with_entrypoint(\"human_converse\")\n",
    "    # add a hook to print the steps -- optional but shows that Burr is pluggable\n",
    "    .with_hooks(PrintStepHook())\n",
    "    # add tracking -- this will show up in the BURR UI.\n",
    "    .with_tracker(tracker)\n",
    "    # build the application\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "34140c5864b940dc",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T20:54:24.035153Z",
     "start_time": "2024-03-26T20:53:19.237522Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# We can now change test, debug, etc. from this prior state.\n",
    "print(f\"Running RAG with loaded state:\\n {pprint.pformat(state_values)}\")\n",
    "while True:\n",
    "    user_question = input(\"Ask something (or type exit to quit): \")\n",
    "    previous_action, result, state = other_app.run(\n",
    "        halt_before=[\"human_converse\"],\n",
    "        halt_after=[\"terminal\"],\n",
    "        inputs={\"user_question\": user_question},\n",
    "    )\n",
    "    if previous_action and previous_action.name == \"terminal\":\n",
    "        # reached the end\n",
    "        pprint.pprint(result)\n",
    "        break\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d7979bc-a09b-4abd-b22b-02ca38d4a26c",
   "metadata": {},
   "source": [
    "### Picking up where we left off from a prior conversation with `initialize_from()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fc62a033644c7b80",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T21:04:56.748649Z",
     "start_time": "2024-03-26T21:04:56.742019Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Now let's show how to use the ApplicationBuilder.initialize_from() method to pick up where we left off.\n",
    "# This is useful if you want to continue a conversation with a user, or if you had a crash, etc.\n",
    "\n",
    "# set up for rewinding to a prior state -- loading it in as initial state\n",
    "prior_app_id = app_id\n",
    "\n",
    "project_name = \"demo:conversational-rag\"\n",
    "# we use the tracking client here to get the state of the application at a prior sequence_id\n",
    "tracker = LocalTrackingClient(project=project_name)\n",
    "pick_up_where_we_left_off_app = (\n",
    "    ApplicationBuilder()\n",
    "    # add the actions\n",
    "    .with_actions(\n",
    "        # bind the vector store to the AI conversational step\n",
    "        ai_converse=ai_converse.bind(vector_store=vector_store),\n",
    "        human_converse=human_converse,\n",
    "        terminal=burr.core.Result(\"chat_history\"),\n",
    "    )\n",
    "    # set the transitions between actions\n",
    "    .with_transitions(\n",
    "        (\"ai_converse\", \"human_converse\", default),\n",
    "        (\"human_converse\", \"terminal\", expr(\"'exit' in question\")),\n",
    "        (\"human_converse\", \"ai_converse\", default),\n",
    "    )\n",
    "    # add identifiers that will help track the application\n",
    "    .with_identifiers(app_id=prior_app_id, partition_key=\"sample_user\")\n",
    "    .initialize_from(\n",
    "        initializer=tracker,\n",
    "        resume_at_next_action=False, # we want to always start at human_converse; our entrypoint\n",
    "        default_entrypoint=\"human_converse\",\n",
    "        default_state=initial_state, # set some default state incase we can't find the prior state\n",
    "    )\n",
    "    # add a hook to print the steps -- optional but shows that Burr is pluggable\n",
    "    .with_hooks(PrintStepHook())\n",
    "    # add tracking -- this will show up in the BURR UI.\n",
    "    .with_tracker(tracker)\n",
    "    # build the application\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b6d23d6d6a6643d0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-26T21:05:41.246005Z",
     "start_time": "2024-03-26T21:05:02.855430Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "print(f\"Running RAG with loaded state:\\n {pprint.pformat(pick_up_where_we_left_off_app.state.get_all())}\")\n",
    "while True:\n",
    "    user_question = input(\"Ask something (or type exit to quit): \")\n",
    "    previous_action, result, state = pick_up_where_we_left_off_app.run(\n",
    "        halt_before=[\"human_converse\"],\n",
    "        halt_after=[\"terminal\"],\n",
    "        inputs={\"user_question\": user_question},\n",
    "    )\n",
    "    if previous_action and previous_action.name == \"terminal\":\n",
    "        # reached the end\n",
    "        pprint.pprint(result)\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e51159c1-04a1-4b55-8615-97dc6d5c6975",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "### Forking the application using prior state with `initialize_from()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "84d533ba-fd85-4b67-a6c3-30356cc9b048",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now let's show how to use the ApplicationBuilder.initialize_from() method to fork a new application.\n",
    "# This is useful if you want to debug/reuse prior state, rewind and try something else, etc.\n",
    "\n",
    "# set up for rewinding to a prior state -- loading it in as initial state\n",
    "prior_app_id = app_id\n",
    "new_app_id = str(uuid.uuid4())\n",
    "\n",
    "project_name = \"demo:conversational-rag\"\n",
    "# we use the tracking client here to get the state of the application at a prior sequence_id\n",
    "tracker = LocalTrackingClient(project=project_name)\n",
    "forking_new_application = (\n",
    "    ApplicationBuilder()\n",
    "    # add the actions\n",
    "    .with_actions(\n",
    "        # bind the vector store to the AI conversational step\n",
    "        ai_converse=ai_converse.bind(vector_store=vector_store),\n",
    "        human_converse=human_converse,\n",
    "        terminal=burr.core.Result(\"chat_history\"),\n",
    "    )\n",
    "    # set the transitions between actions\n",
    "    .with_transitions(\n",
    "        (\"ai_converse\", \"human_converse\", default),\n",
    "        (\"human_converse\", \"terminal\", expr(\"'exit' in question\")),\n",
    "        (\"human_converse\", \"ai_converse\", default),\n",
    "    )\n",
    "    # add identifiers that will help track the application\n",
    "    .with_identifiers(app_id=new_app_id, partition_key=\"sample_user\")\n",
    "    .initialize_from(\n",
    "        initializer=tracker,\n",
    "        resume_at_next_action=False, # we want to always start at human_converse; our entrypoint\n",
    "        default_entrypoint=\"human_converse\",\n",
    "        default_state=initial_state, # set some default state incase we can't find the prior state\n",
    "        fork_from_app_id=prior_app_id, # <---- The new addition \n",
    "        fork_from_partition_key=\"sample_user\",  # <---- The new addition \n",
    "        fork_from_sequence_id=5,  # <---- The new addition \n",
    "    )\n",
    "    # add a hook to print the steps -- optional but shows that Burr is pluggable\n",
    "    .with_hooks(PrintStepHook())\n",
    "    # add tracking -- this will show up in the BURR UI.\n",
    "    .with_tracker(tracker)\n",
    "    # build the application\n",
    "    .build()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a9086e58-eeaa-4886-a113-5c7fd5fe0af0",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Running RAG with loaded state:\\n {pprint.pformat(forking_new_application.state.get_all())}\")\n",
    "while True:\n",
    "    user_question = input(\"Ask something (or type exit to quit): \")\n",
    "    previous_action, result, state = forking_new_application.run(\n",
    "        halt_before=[\"human_converse\"],\n",
    "        halt_after=[\"terminal\"],\n",
    "        inputs={\"user_question\": user_question},\n",
    "    )\n",
    "    if previous_action and previous_action.name == \"terminal\":\n",
    "        # reached the end\n",
    "        pprint.pprint(result)\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d6523a54-e333-4d5e-941d-d3b28b940500",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this should print some warnings:\n",
    "new_app_id = str(uuid.uuid4())\n",
    "\n",
    "project_name = \"demo:conversational-rag\"\n",
    "# we use the tracking client here to get the state of the application at a prior sequence_id\n",
    "tracker = LocalTrackingClient(project=project_name)\n",
    "bad_fork = (\n",
    "    ApplicationBuilder()\n",
    "    # add the actions\n",
    "    .with_actions(\n",
    "        # bind the vector store to the AI conversational step\n",
    "        ai_converse=ai_converse.bind(vector_store=vector_store),\n",
    "        human_converse=human_converse,\n",
    "        terminal=burr.core.Result(\"chat_history\"),\n",
    "    )\n",
    "    # set the transitions between actions\n",
    "    .with_transitions(\n",
    "        (\"ai_converse\", \"human_converse\", default),\n",
    "        (\"human_converse\", \"terminal\", expr(\"'exit' in question\")),\n",
    "        (\"human_converse\", \"ai_converse\", default),\n",
    "    )\n",
    "    # add identifiers that will help track the application\n",
    "    .with_identifiers(app_id=new_app_id, partition_key=\"sample_user\")\n",
    "    .initialize_from(\n",
    "        initializer=tracker,\n",
    "        resume_at_next_action=False, # we want to always start at human_converse; our entrypoint\n",
    "        default_entrypoint=\"human_converse\",\n",
    "        default_state=initial_state, # set some default state incase we can't find the prior state\n",
    "        fork_from_app_id=\"NO_SUCH_APP_ID\", # <---- The new addition \n",
    "    )\n",
    "    # add a hook to print the steps -- optional but shows that Burr is pluggable\n",
    "    .with_hooks(PrintStepHook())\n",
    "    # add tracking -- this will show up in the BURR UI.\n",
    "    .with_tracker(tracker)\n",
    "    # build the application\n",
    "    .build()\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
